from jqdata import *
import pandas as pd

# 利用帮助文档表格，直接读取指标及名称对应的字典
balance_dict=pd.read_excel("单季财务数据字典.xlsx",sheet_name="Balance")
cashflow_dict=pd.read_excel("单季财务数据字典.xlsx",sheet_name="Cashflow")
income_dict=pd.read_excel("单季财务数据字典.xlsx",sheet_name="Income")
balance_dict["列的含义"]=balance_dict["列的含义"].str.strip("(元)")
cashflow_dict["列的含义"]=cashflow_dict["列的含义"].str.strip("(元)")
income_dict["列的含义"]=income_dict["列的含义"].str.strip("(元)")

def find_indicator(indicator):
    """财务指标查询"""
    for mapstr in ["balance","cashflow","income"]:
        maps=globals()[mapstr+"_dict"]
        if indicator in maps["列名"].values:
            indicator_cn=maps.loc[maps["列名"]==indicator,"列的含义"].iloc[0]
            return f"{mapstr}.{indicator}",mapstr,indicator_cn
        elif indicator in maps["列的含义"].values:
            indicator_cn=indicator
            indicator=maps.loc[maps["列的含义"]==indicator,"列名"].iloc[0]
            return f"{mapstr}.{indicator}",mapstr,indicator_cn
    raise Exception("没有找到指标!")


def mydataq(code,indicator,start_quarter,end_quarter):
    """
    :参数 code(股票代码),indicator(财务指标中文名),
            start_quarter,end_quarter(季度区间)
    :作用 查询季度财务指标
    :返回 data(给定时间段内、指定公司的财务指标)
    """
    year,quarter_num=start_quarter.split("q")
    year=int(year)
    quarter_num=int(quarter_num)
    quarter_=start_quarter
    quarter_series=[start_quarter]
    while quarter_!=end_quarter:
        if quarter_num==4:
            year+=1
            quarter_num=1
        else:
            quarter_num+=1
        quarter_=f"{year}q{quarter_num}"
        quarter_series.append(quarter_)
    expression,table,cn_name=find_indicator(indicator)
    q=query(eval(expression)).filter(eval(table).code==code)
    effected_quarter=[]
    data=[]
    for x in quarter_series:
        single=get_fundamentals(q,statDate=x).dropna()
        if single.empty:
            print(f"{code}在{x}无数据")
        else:
            data.append(single.iat[0,0])
            effected_quarter.append(x)
    data=pd.Series(data,index=effected_quarter).to_frame()
    data.rename(columns={0:indicator},inplace=True)
    return data

